#!/usr/bin/env Rscript
# -*- coding: utf-8 -*-
# ------------------------------------------------------------------------------
# Beta 多样性分析：Bray-Curtis 距离 + NMDS + PCoA + 相关性热图（带分组信息）
# 使用 ComplexHeatmap 替代 pheatmap 绘制相关性热图
# ------------------------------------------------------------------------------

suppressPackageStartupMessages({
  library(optparse)
  library(tidyverse)
  library(vegan)
  library(ggplot2)
  library(ComplexHeatmap)
  library(circlize)
  library(RColorBrewer)
  library(ggpubr)
  library(rstatix)
})

# ----------------------
# 日志打印
# ----------------------
log.message <- function(msg) {
  message(sprintf("[%s] %s", Sys.time(), msg))
}

# ----------------------
# 参数解析
# ----------------------
parse.opt <- function() {
  option_list <- list(
    make_option(c("-i", "--input"), type="character", help="RPM矩阵（行为物种，列为样本）"),
    make_option(c("-m", "--meta"), type="character", help="分组文件（两列：Sample, Group）"),
    make_option(c("-o", "--output"), type="character", help="Beta距离矩阵输出"),
    make_option(c("--nmds_pdf"), type="character", help="NMDS PDF"),
    make_option(c("--nmds_png"), type="character", help="NMDS PNG"),
    make_option(c("--pcoa_pdf"), type="character", help="PCoA PDF"),
    make_option(c("--pcoa_png"), type="character", help="PCoA PNG"),
    make_option(c("--corr_pdf"), type = "character", help="相关性热图 PDF"),
    make_option(c("--corr_png"), type="character", help="相关性热图 PNG"),
    make_option(c("--spec_pdf"), type = "character", help="物种丰富热图PDF"),
    make_option(c("--spec_png"), type = "character", help = "物种丰富热图PNG"),
    make_option(c("--width"), type="numeric", default=8, help="图宽度"),
    make_option(c("--height"), type="numeric", default=6, help="图高度"),
    make_option(c("--dpi"), type="numeric", default=300, help="PNG 分辨率")
  )
  parser <- OptionParser(option_list=option_list)
  opt <- parse_args(parser)
  
  required <- c("input","meta","output","nmds_pdf","nmds_png","pcoa_pdf",
                "pcoa_png","corr_pdf","corr_png", "spec_pdf", "spec_png")
  miss <- required[!required %in% names(opt) | sapply(opt[required], is.null)]
  if(length(miss) > 0) {
    print_help(parser)
    stop(paste("缺少必填参数:", paste(miss, collapse=", ")))
  }
  return(opt)
}

# ----------------------
# 数据读取
# ----------------------
read.input <- function(path) {
  df <- read.table(path, header=TRUE, sep="\t", row.names=1, check.names=FALSE)
  return(t(df))  # 转置：样本为行
}

read.meta <- function(path) {
  df <- read.table(path, header=TRUE, sep="\t", stringsAsFactors=FALSE)
  df$Sample <- trimws(df$Sample)  # 去掉空格
  return(df)
}

# ----------------------
# NMDS 绘图
# ----------------------
nmds.plot <- function(beta_div, meta_df, out.pdf, out.png, width, height, dpi) {
  # 增加 trymax 减少收敛警告
  nmds.res <- metaMDS(beta_div, k=2, trymax=2000)
  points <- as.data.frame(nmds.res$points)
  points$Sample <- rownames(points)
  points <- left_join(points, meta_df, by="Sample")
  
  p <- ggplot(points, aes(x=MDS1, y=MDS2, color=Group)) +
    geom_point(size=3) +
    stat_ellipse(aes(fill=Group), geom="polygon", alpha=0.2, color=NA, type="t", level=0.95) +
    theme_bw() +
    labs(title=paste("NMDS (", attr(beta_div,"method"), ")", sep=""))
  
  ggsave(out.pdf, plot=p, width=width, height=height)
  ggsave(out.png, plot=p, width=width, height=height, dpi=dpi)
}

# ----------------------
# PCoA 绘图
# ----------------------
pcoa.plot <- function(beta_div, meta_df, out.pdf, out.png, width, height, dpi) {
  pcoa.res <- cmdscale(beta_div, eig=TRUE, k=2)
  points <- as.data.frame(pcoa.res$points)
  colnames(points) <- c("PCoA1","PCoA2")
  points$Sample <- rownames(points)
  points <- left_join(points, meta_df, by="Sample")
  
  p <- ggplot(points, aes(x=PCoA1, y=PCoA2, color=Group)) +
    geom_point(size=3) +
    stat_ellipse(aes(fill=Group), geom="polygon", alpha=0.2, color=NA, type="t", level=0.95) +
    theme_bw() +
    labs(title=paste("PCoA (", attr(beta_div,"method"), ")", sep=""))
  
  ggsave(out.pdf, plot=p, width=width, height=height)
  ggsave(out.png, plot=p, width=width, height=height, dpi=dpi)
}

# ----------------------
# Beta 距离热图（ComplexHeatmap）
# ----------------------
beta.heatmap <- function(beta_div, meta_df, out.pdf, out.png, width=8, height=6, dpi=360) {
  
  # 转成矩阵（beta_div 是 dist 对象）
  beta_mat <- as.matrix(beta_div)
  
  # -------------------------
  # 对齐 meta_df 与 beta_mat 样本
  # -------------------------
  meta_df$Group <- as.factor(meta_df$Group)
  
  # 找出共同的样本
  common_samples <- intersect(rownames(beta_mat), meta_df$Sample)
  if(length(common_samples) == 0) {
    stop("❌ 没有匹配的样本，请检查 meta_df$Sample 与 beta_div 行名是否一致！")
  }
  
  beta_mat <- beta_mat[common_samples, common_samples]
  meta_df <- meta_df[meta_df$Sample %in% common_samples, ]
  meta_df <- meta_df[match(common_samples, meta_df$Sample), ]
  
  # -------------------------
  # 分组颜色
  # -------------------------
  group.colors <- setNames(
    if(length(levels(meta_df$Group)) > 1) {
      colorRampPalette(RColorBrewer::brewer.pal(8, "Set1"))(length(levels(meta_df$Group)))
    } else {
      "#E41A1C"
    },
    levels(meta_df$Group)
  )
  
  # -------------------------
  # 行列注释
  # -------------------------
  col_ha <- HeatmapAnnotation(
    Group = meta_df$Group,
    col = list(Group = group.colors),
    annotation_name_side = "right"
  )
  
  row_ha <- rowAnnotation(
    Group = meta_df$Group,
    col = list(Group = group.colors),
    annotation_name_side = "top",
    show_legend = FALSE
  )
  
  # -------------------------
  # 绘制热图
  # -------------------------
  ht <- Heatmap(
    beta_mat,
    name = "Bray-Curtis",
    top_annotation = col_ha,
    left_annotation = row_ha,
    show_row_names = FALSE,
    show_column_names = FALSE,
    clustering_method_rows = "average",
    clustering_method_columns = "average",
    col = colorRamp2(c(min(beta_mat), max(beta_mat)), c("white", "red"))  # 距离越大颜色越深
  )
  
  pdf(out.pdf, width=width, height=height)
  draw(ht)
  dev.off()
  
  png(out.png, width=width, height=height, units="in", res=dpi)
  draw(ht)
  dev.off()
}

# ----------------------
# 物种丰度聚类热图（ComplexHeatmap） - 带诊断信息与更健壮处理
# ----------------------
species.heatmap <- function(input_path, meta_df, out.pdf, out.png, width=10, height=8, dpi=360) {
  log.message(sprintf("开始绘制物种丰度聚类热图: input=%s, out.pdf=%s, out.png=%s", input_path, out.pdf, out.png))
  
  if (!file.exists(input_path)) {
    warning("输入文件不存在: ", input_path)
    return(invisible(NULL))
  }
  
  # 读取物种丰度矩阵（行=species, 列=samples）
  mat <- tryCatch(
    read.table(input_path, header=TRUE, sep="\t", row.names=1, check.names=FALSE),
    error = function(e) {
      warning("读取物种丰度矩阵失败: ", e$message)
      return(NULL)
    }
  )
  if (is.null(mat)) return(invisible(NULL))
  
  log.message(sprintf("原始矩阵大小: %d species x %d samples", nrow(mat), ncol(mat)))
  
  # 去掉全零行和列
  nz_row <- rowSums(mat) > 0
  nz_col <- colSums(mat) > 0
  mat <- mat[nz_row, nz_col, drop=FALSE]
  log.message(sprintf("过滤后矩阵大小: %d species x %d samples", nrow(mat), ncol(mat)))
  
  if (nrow(mat) < 2) {
    warning("物种数量 < 2，跳过物种热图绘制。")
    return(invisible(NULL))
  }
  if (ncol(mat) < 2) {
    warning("样本数量 < 2，跳过物种热图绘制。")
    return(invisible(NULL))
  }
  
  # 对齐样本顺序，注意清理空格并强制字符
  meta_df$Sample <- as.character(trimws(meta_df$Sample))
  meta_df$Group  <- as.factor(meta_df$Group)
  
  common_samples <- intersect(colnames(mat), meta_df$Sample)
  log.message(sprintf("匹配到的共同样本数: %d", length(common_samples)))
  if (length(common_samples) == 0) {
    # 打印提示，列出前若干个样本名做比对方便排查
    msg <- paste0(
      "没有共同样本 (colnames(mat) 与 meta_df$Sample 不匹配).\n",
      "mat samples (first 10): ", paste(head(colnames(mat), 10), collapse=", "), "\n",
      "meta samples (first 10): ", paste(head(meta_df$Sample, 10), collapse=", "), "\n",
      "请检查样本名是否包含空格、特殊字符或顺序差异。"
    )
    stop(msg)
  }
  
  # 子集化并保持 meta_df 顺序与样本一致
  mat <- mat[, common_samples, drop=FALSE]
  meta_df <- meta_df[match(common_samples, meta_df$Sample), , drop=FALSE]
  
  # 对数化 / 标准化（按物种行）
  mat_log <- log10(mat + 1)
  # 当某些物种方差为 0 时 scale 会产生 NA，需要处理
  row_sds <- apply(mat_log, 1, sd, na.rm=TRUE)
  if (any(row_sds == 0 | is.na(row_sds))) {
    # 对于方差为0的行，直接置0（即标准化后为0）
    mat_scaled <- mat_log
    zero_rows <- which(row_sds == 0 | is.na(row_sds))
    if (length(zero_rows) > 0) {
      mat_scaled[zero_rows, ] <- 0
    }
    # 其他行标准化
    non_zero_rows <- setdiff(seq_len(nrow(mat_log)), zero_rows)
    if (length(non_zero_rows) > 0) {
      mat_scaled[non_zero_rows, ] <- t(scale(t(mat_log[non_zero_rows, , drop=FALSE])))
    }
  } else {
    mat_scaled <- t(scale(t(mat_log)))
  }
  
  # 分组颜色
  group_colors <- if (length(levels(meta_df$Group)) > 1) {
    setNames(colorRampPalette(RColorBrewer::brewer.pal(8, "Set1"))(length(levels(meta_df$Group))),
             levels(meta_df$Group))
  } else {
    setNames("#E41A1C", levels(meta_df$Group))
  }
  
  top_ha <- HeatmapAnnotation(
    Group = meta_df$Group,
    col = list(Group = group_colors),
    annotation_name_side = "right"
  )
  
  # 根据样本数决定是否显示列名（样本太多会挤在一起）
  show_colnames_flag <- ncol(mat_scaled) <= 60
  
  ht <- Heatmap(
    mat_scaled,
    name = "log10(RPM+1)",
    top_annotation = top_ha,
    show_row_names = FALSE,
    show_column_names = show_colnames_flag,
    column_names_rot = ifelse(show_colnames_flag, 45, 0),
    clustering_method_rows = "average",
    clustering_method_columns = "average",
    col = colorRamp2(c(-2, 0, 2), c("navy", "white", "firebrick"))
  )
  
  # 确保输出目录存在
  outdir_pdf <- dirname(out.pdf)
  outdir_png <- dirname(out.png)
  if (!dir.exists(outdir_pdf)) dir.create(outdir_pdf, recursive = TRUE, showWarnings = FALSE)
  if (!dir.exists(outdir_png)) dir.create(outdir_png, recursive = TRUE, showWarnings = FALSE)
  
  # 写文件并捕获错误
  tryCatch({
    pdf(out.pdf, width=width, height=height)
    draw(ht)
    dev.off()
    log.message(sprintf("已写入 PDF: %s", out.pdf))
  }, error = function(e) {
    warning("写 PDF 失败: ", e$message)
  })
  
  tryCatch({
    png(out.png, width=width, height=height, units="in", res=dpi)
    draw(ht)
    dev.off()
    log.message(sprintf("已写入 PNG: %s", out.png))
  }, error = function(e) {
    warning("写 PNG 失败: ", e$message)
  })
  
  invisible(NULL)
}

# ----------------------
# 主函数
# ----------------------
main <- function() {
  opt <- parse.opt()
  log.message("读取数据...")
  data_t <- read.input(opt$input)
  meta_df <- read.meta(opt$meta)
  
  log.message("计算 Beta 多样性...")
  data_t <- data_t[rowSums(data_t) > 0, ]
  data_t <- data_t[, colSums(data_t) > 0]
  beta_div <- vegdist(data_t, method="bray")
  
  log.message("绘制物种丰度聚类热图...")
  species.heatmap(opt$input, 
                  meta_df,
                  opt$spec_pdf,
                  opt$spec_png,
                  width = opt$width * 1.2,
                  height = opt$height * 1.2,
                  dpi = opt$dpi)
  
  
  log.message("保存 Beta 距离矩阵...")
  dir.create(dirname(opt$output), recursive=TRUE, showWarnings=FALSE)
  write.table(as.matrix(beta_div), file=opt$output, sep="\t", quote=FALSE)
  
  log.message("绘制相关性热图...")
  beta.heatmap(beta_div, meta_df, opt$corr_pdf, opt$corr_png)
  
  log.message("绘制 NMDS...")
  nmds.plot(beta_div, meta_df, opt$nmds_pdf, opt$nmds_png, opt$width, opt$height, opt$dpi)
  
  log.message("绘制 PCoA...")
  pcoa.plot(beta_div, meta_df, opt$pcoa_pdf, opt$pcoa_png, opt$width, opt$height, opt$dpi)
  
  log.message("分析完成 ✅")
}

main()
